import paramiko
import re
import json

# 示例业务逻辑
def example_logic():
    # 返回一个简单的JSON响应数据
    return {"message": "Hello, this is an example API!"}

# 新增：SSH连接服务器
def ssh_connect(hostname, port, username, password):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(hostname, port=port, username=username, password=password)
    return client

# 修改：SSH执行命令
def ssh_execute_command(client, command):
    stdin, stdout, stderr = client.exec_command(command)
    output = stdout.read().decode()
    error = stderr.read().decode()
    client.close()
    return {"output": output, "error": error}

# 新增：提取设备信息
def extract_device_info(device_list, base_list):
    dev_list = []
    base_list_res = []

    for device in device_list:
        device_info = {
            "deviceType": device.get("deviceType"),
            "manufacturer": device.get("manufacturer"),
            "model": device.get("model"),
            "firmwareVersion": device.get("firmwareVersion"),
            "ruleVersion": None,  # 需要从策略中提取
            "remark": "",
            "strategies": analyze_strategies(device.get("strategies", []))
        }
        dev_list.append(device_info)

    for base in base_list:
        base_info = {
            "deviceType": base.get("deviceType"),
            "manufacturer": base.get("manufacturer"),
            "model": base.get("model"),
            "firmwareVersion": base.get("firmwareVersion"),
            "ruleVersion": None,  # 需要从策略中提取
            "remark": base.get("remark", "")
        }
        base_list_res.append(base_info)

    return dev_list, base_list_res

# 新增：比较版本号
def compare_versions(version1, version2):
    def parse_version(version):
        return [int(part) if part.isdigit() else part for part in re.split(r'[.-]', version)]

    v1_parts = parse_version(version1)
    v2_parts = parse_version(version2)

    for v1, v2 in zip(v1_parts, v2_parts):
        if v1 < v2:
            return -1
        elif v1 > v2:
            return 1

    if len(v1_parts) < len(v2_parts):
        return -1
    elif len(v1_parts) > len(v2_parts):
        return 1

    return 0

# 新增：分析策略
def analyze_strategies(strategies):
    analysis_results = []
    for strategy in strategies:
        info = strategy.get("info", {})
        analysis = []

        # 示例分析逻辑
        if "rules" in info:
            for rule in info["rules"]:
                if rule["action"] == "allow" and rule["src_zone"] == "untrust" and rule["service"] == "tcp/3389" and not rule["log"]:
                    analysis.append("High risk: Allow untrusted access to RDP without logging")
                # 其他分析逻辑...

        strategy["analysis"] = ", ".join(analysis)
        analysis_results.append(strategy)

    return analysis_results

# 新增：测试方法
def test_ssh_connect():
    # 测试SSH连接
    client = ssh_connect('localhost', 22, 'username', 'password')
    assert client is not None
    client.close()

def test_ssh_execute_command():
    # 测试SSH执行命令
    client = ssh_connect('localhost', 22, 'username', 'password')
    result = ssh_execute_command(client, 'ls -l')
    assert 'total' in result['output']
    assert result['error'] == ''

def test_extract_device_info():
    # 测试提取设备信息
    device_list = [
        {
            "deviceType": "Router",
            "manufacturer": "Cisco",
            "model": "ISR 4461",
            "firmwareVersion": "16.12.2",
            "strategies": [
                {
                    "info": {
                        "rules": [
                            {
                                "action": "allow",
                                "src_zone": "untrust",
                                "service": "tcp/3389",
                                "log": False
                            }
                        ]
                    }
                }
            ]
        }
    ]
    base_list = [
        {
            "deviceType": "Router",
            "manufacturer": "Cisco",
            "model": "ISR 4461",
            "firmwareVersion": "16.12.2",
            "remark": "Base device"
        }
    ]
    dev_list, base_list_res = extract_device_info(device_list, base_list)
    assert dev_list[0]['strategies'][0]['analysis'] == "High risk: Allow untrusted access to RDP without logging"
    assert base_list_res[0]['remark'] == "Base device"

def test_compare_versions():
    # 测试比较版本号
    assert compare_versions('1.0.0', '1.0.1') == -1
    assert compare_versions('1.0.1', '1.0.0') == 1
    assert compare_versions('1.0.0', '1.0.0') == 0
    assert compare_versions('1.0', '1.0.0') == 0
    assert compare_versions('1.0.0', '1.0') == 0

def test_analyze_strategies():
    # 测试分析策略
    strategies = [
        {
            "info": {
                "rules": [
                    {
                        "action": "allow",
                        "src_zone": "untrust",
                        "service": "tcp/3389",
                        "log": False
                    }
                ]
            }
        }
    ]
    analysis_results = analyze_strategies(strategies)
    assert analysis_results[0]['analysis'] == "High risk: Allow untrusted access to RDP without logging"

def parse_version(version):
    """支持处理含字母、数字的混合版本格式（如6R10F02）"""
    version_str = str(version).lower().strip('v')
    # 使用正则提取所有数字段（支持R/F等字母分隔符）
    parts = re.findall(r'\d+', version_str)
    return [int(p) for p in parts]

def analyze_strategy(info):
    try:
        info_dict = json.loads(info.replace('{{{{', '{{').replace('}}}}', '}}'))
        rules = info_dict.get('rules', [])
        risk_level = '低风险'
        optimization = '无'
        tightening = '无'

        for rule in rules:
            if rule.get('action') == 'allow' and rule.get('src_zone') == 'untrust' and rule.get(
                    'src_ip') == 'any' and rule.get('service') == 'tcp/3389' and not rule.get('log', False):
                risk_level = '高风险'
            # 这里可以继续添加其他规则判断优化和收紧情况
        return f"策略风险：{risk_level}，策略优化：{optimization}，策略收紧：{tightening}"
    except Exception:
        return "策略信息解析失败"

def process_data(input_data):
    # 若 input_data 是字符串，解析为字典
    if isinstance(input_data, str):
        input_data = json.loads(input_data)

    device_list = input_data.get('deviceList', [])
    base_list = input_data.get('baseList', [])
    device_result_list = []
    base_result_list = []

    for device in device_list:
        device_info = {
            "deviceType": device.get('deviceType'),
            "manufacturer": device.get('manufacturer'),
            "model": device.get('model'),
            "firmwareVersion": device.get('firmwareVersion'),
            "ruleVersion": None,
            "strategies": []
        }
        strategies = device.get('strategies', [])
        for strategy in strategies:
            analysis = analyze_strategy(strategy.get('info', '{}'))
            device_info["strategies"].append({
                "id": strategy.get('id'),
                "name": strategy.get('name'),
                "type": strategy.get('type'),
                "analysis": analysis
            })
        for base_device in base_list:
            if (device.get('deviceType') == base_device.get('deviceType') and
                    device.get('manufacturer') == base_device.get('manufacturer') and
                    device.get('model') == base_device.get('model')):
                firmware_compare = compare_versions(device.get('firmwareVersion'),
                                                    base_device.get('firmwareVersion'))
                if firmware_compare < 0:
                    device_info[
                        "remark"] = f"软件/固件版本不合规，最新版本号为 {base_device.get('firmwareVersion')}  ，分析思路：当前设备软件/固件版本低于基线版本"
                elif firmware_compare > 0:
                    base_result_list.append({
                        "deviceType": base_device.get('deviceType'),
                        "manufacturer": base_device.get('manufacturer'),
                        "model": base_device.get('model'),
                        "firmwareVersion": base_device.get('firmwareVersion'),
                        "ruleVersion": None,
                        "remark": f"基线数据过时，最新版本号为 {device.get('firmwareVersion')}  ，分析思路：设备中的软件/固件版本高于基线版本"
                    })
        if "remark" in device_info or device_info["strategies"]:
            device_result_list.append(device_info)

    return {
        "devList": device_result_list,
        "baseListRes": base_result_list
    }